1
За пределами элементарных операций: переход к тайловым матричным операциям
AI023Lesson 9
00:00

В предыдущих уроках мы сосредоточились на элементарных операциях (например, простой функции ReLU для матрицы). Эти операции являются ограниченными памятью потому что видеопроцессор тратит больше времени на перемещение данных из памяти HBM в регистры, чем на выполнение вычислений.

1. Почему GEMM является центральным

Общее умножение матриц (GEMM) имеет вычислительную сложность $O(N^3)$, при этом требуя лишь $O(N^2)$ доступа к памяти. Это позволяет скрывать задержки памяти за огромной производительностью арифметических операций, делая его «сердцем» моделей больших языковых моделей (LLM).

2. Представление памяти в двумерном виде

Физическая ОЗУ — одномерная. Чтобы представить двумерный тензор, мы используем шаги (strides). Распространённая ошибка в производственной среде — предположение, что тензор континуален. Если вы перепутаете шаги строк и столбцов в вашем коде указателей, вы получите доступ к «призрачным» данным или вызовете нарушения памяти.

3. Обобщение с использованием тайлов

Triton обобщает логику элементарных операций, переходя от одиночных указателей к блокам указателей. Используя двумерные тайлы (например, $16 \times 16$), мы используем восстановление данных в высокоскоростной памяти SRAM, сохраняя данные «горячими» для объединённых операций, таких как добавление смещения (bias) или активации, перед записью обратно в глобальную память.

Линейное представление 1DТайловое представление 2D
main.py
TERMINALbash — 80x24
> Ready. Click "Run" to execute.
>